import numpy as np
import pandas as pd
import talib
import matplotlib.pyplot as plt

# 假设我们有1000组数据，这里生成示例数据
np.random.seed(42)
n_samples = 1000

# 生成模拟价格数据 (开盘价、收盘价、最高价、最低价)
close_prices = np.cumprod(1 + np.random.normal(0.001, 0.02, n_samples)) * 100
open_prices = close_prices * (1 + np.random.normal(0, 0.005, n_samples))
high_prices = close_prices * (1 + np.abs(np.random.normal(0, 0.01, n_samples)))
low_prices = close_prices * (1 - np.abs(np.random.normal(0, 0.01, n_samples)))

# 生成模拟成交量数据
volumes = np.random.randint(10000, 50000, n_samples)

# 创建DataFrame
data = pd.DataFrame({
    'open': open_prices,
    'high': high_prices,
    'low': low_prices,
    'close': close_prices,
    'volume': volumes
})


def find_support_resistance_levels(data, price_range_percent=[0.01, 0.03, 0.07]):
    """
    寻找密集成交区的支撑位和阻力位

    参数:
    data: 包含OHLC和成交量数据的DataFrame
    price_range_percent: 强度级别的价格区间百分比(弱小,中等,高强度)

    返回:
    dict: 包含支撑位和阻力位的字典，按强度分类
    """
    close_prices = data['close'].values
    volumes = data['volume'].values

    # 计算价格范围
    price_range = np.max(close_prices) - np.min(close_prices)

    # 定义强度级别
    strength_levels = {
        'weak': price_range_percent[0],
        'medium': price_range_percent[1],
        'strong': price_range_percent[2]
    }

    # 计算价格直方图 (成交量加权)
    hist, bin_edges = np.histogram(close_prices, bins=50, weights=volumes)

    # 找到直方图中的峰值 (高成交量区域)
    peaks = []
    for i in range(1, len(hist) - 1):
        if hist[i] > hist[i - 1] and hist[i] > hist[i + 1]:
            peaks.append((bin_edges[i], hist[i]))

    # 按成交量排序峰值
    peaks_sorted = sorted(peaks, key=lambda x: x[1], reverse=True)

    # 分类支撑/阻力位
    levels = {'support': {'weak': [], 'medium': [], 'strong': []},
              'resistance': {'weak': [], 'medium': [], 'strong': []}}

    current_price = close_prices[-1]

    for price, vol in peaks_sorted:
        # 确定是支撑还是阻力
        level_type = 'support' if price < current_price else 'resistance'

        # 计算与当前价格的距离百分比
        distance_pct = abs(price - current_price) / current_price

        # 确定强度级别
        if distance_pct <= strength_levels['weak']:
            strength = 'weak'
        elif distance_pct <= strength_levels['medium']:
            strength = 'medium'
        elif distance_pct <= strength_levels['strong']:
            strength = 'strong'
        else:
            continue  # 忽略太远的水平

        levels[level_type][strength].append(price)

    return levels


def plot_levels(data, levels):
    """
    绘制价格走势和支撑/阻力位
    """
    plt.figure(figsize=(12, 6))
    plt.plot(data['close'], label='Price')

    # 绘制支撑位
    for strength, color in zip(['weak', 'medium', 'strong'], ['green', 'blue', 'red']):
        for level in levels['support'][strength]:
            plt.axhline(y=level, color=color, linestyle='--', alpha=0.5,
                        label=f'Support ({strength})' if strength == 'weak' else "")

    # 绘制阻力位
    for strength, color in zip(['weak', 'medium', 'strong'], ['green', 'blue', 'red']):
        for level in levels['resistance'][strength]:
            plt.axhline(y=level, color=color, linestyle='-', alpha=0.5,
                        label=f'Resistance ({strength})' if strength == 'weak' else "")

    plt.title('Support and Resistance Levels')
    plt.xlabel('Time')
    plt.ylabel('Price')
    plt.legend()
    plt.show()


# 寻找支撑阻力位
levels = find_support_resistance_levels(data)

# 打印结果
print("支撑位:")
print(f"弱小 (1%): {np.round(levels['support']['weak'], 2)}")
print(f"中等 (3%): {np.round(levels['support']['medium'], 2)}")
print(f"高强度 (7%): {np.round(levels['support']['strong'], 2)}")

print("\n阻力位:")
print(f"弱小 (1%): {np.round(levels['resistance']['weak'], 2)}")
print(f"中等 (3%): {np.round(levels['resistance']['medium'], 2)}")
print(f"高强度 (7%): {np.round(levels['resistance']['strong'], 2)}")

# 绘制图表
plot_levels(data, levels)